#include <stdio.h>
#include <stdlib.h>
#include <string.h>
// #include <math.h>

#include "darknet_TA.h"
#include "blas_TA.h"
#include "network_TA.h"
#include "math_TA.h"

#include "darknetp_ta.h"
#include <tee_internal_api.h>
#include <tee_internal_api_extensions.h>
#include <tee_api.h>
#include "activations_TA.h"

#include "activation_layer_TA.h"

network_TA netta;
int roundnum = 0;
float err_sum = 0;
float avg_loss = -1;

float *ta_net_input;
float *ta_net_delta;
float *ta_net_output;

void make_network_TA(int n, float learning_rate, float momentum, float decay, int time_steps, int notruth, int batch, int subdivisions, int random, int adam, float B1, float B2, float eps, int h, int w, int c, int inputs, int max_crop, int min_crop, float max_ratio, float min_ratio, int center, float clip, float angle, float aspect, float saturation, float exposure, float hue, int burn_in, float power, int max_batches)
{
    netta.n = n;

    //netta.seen = calloc(1, sizeof(size_t));
    netta.seen = calloc(1, sizeof(uint64_t));
    netta.layers = calloc(netta.n, sizeof(layer_TA));
    netta.t    = calloc(1, sizeof(int));
    netta.cost = calloc(1, sizeof(float));

    netta.learning_rate = learning_rate;
    netta.momentum = momentum;
    netta.decay = decay;
    netta.time_steps = time_steps;
    netta.notruth = notruth;
    netta.batch = batch;
    netta.subdivisions = subdivisions;
    netta.random = random;
    netta.adam = adam;
    netta.B1 = B1;
    netta.B2 = B2;
    netta.eps = eps;
    netta.h = h;
    netta.w = w;
    netta.c = c;
    netta.inputs = inputs;
    netta.max_crop = max_crop;
    netta.min_crop = min_crop;
    netta.max_ratio = max_ratio;
    netta.min_ratio = min_ratio;
    netta.center = center;
    netta.clip = clip;
    netta.angle = angle;
    netta.aspect = aspect;
    netta.saturation = saturation;
    netta.exposure = exposure;
    netta.hue = hue;
    netta.burn_in = burn_in;
    netta.power = power;
    netta.max_batches = max_batches;
    netta.workspace_size = 0;

    //netta.truth = net->truth; ////// ing network.c train_network
}
static inline uint32_t tee_time_to_ms(TEE_Time t)
{
	return t.seconds * 1000 + t.millis;
}

static inline uint32_t get_delta_time_in_ms(TEE_Time start, TEE_Time stop)
{
	return tee_time_to_ms(stop) - tee_time_to_ms(start);
}

void forward_network_TA()
{
    if(roundnum == 0){
        // ta_net_input malloc so not destroy before addition backward
        ta_net_input = malloc(sizeof(float) * netta.layers[0].inputs * netta.layers[0].batch);
        ta_net_delta = malloc(sizeof(float) * netta.layers[0].inputs * netta.layers[0].batch);
        // ta_net_input = malloc(sizeof(float) * netta.layers[0].inputs * netta.layers[0].batch);
        // ta_net_delta = malloc(sizeof(float) * netta.layers[0].inputs * netta.layers[0].batch);

        if(netta.workspace_size){
            printf("workspace_size=%ld\n", netta.workspace_size);
            netta.workspace = calloc(1, netta.workspace_size);
        }
    }

    roundnum++;
    int i;
    TEE_Time start_time = { };
	TEE_Time stop_time = { };
	TEE_GetSystemTime(&start_time);

    for(i = 0; i < netta.n; ++i){
        netta.index = i;
        layer_TA l = netta.layers[i];

        if(l.delta){
            fill_cpu_TA(l.outputs * l.batch, 0, l.delta, 1);
        }

        l.forward_TA(l, netta);

        // if(debug_summary_pass == 1){
        //     summary_array("forward_network / l.output", l.output, l.outputs*netta.batch);
        // }

        netta.input = l.output;

        if(l.truth) {
            netta.truth = l.output;
        }
        //output of the network (for predict)
        // &&
        // if(!netta.train && l.type == SOFTMAX_TA){
        //     ta_net_output = malloc(sizeof(float)*l.outputs*1);
        //     for(int z=0; z<l.outputs*1; z++){
        //         ta_net_output[z] = l.output[z];
        //     }
        // }

        if(i == netta.n - 1)  // ready to back REE for the rest forward pass
        {
            ta_net_input = malloc(sizeof(float)*l.outputs*l.batch);
            for(int z=0; z<l.outputs*l.batch; z++){
                ta_net_input[z] = netta.input[z];
            }
        }
    }
	TEE_GetSystemTime(&stop_time);
    printf("InVoked: start: %u.%u(s), stop: %u.%u(s), delta: %u(ms)",
			start_time.seconds, start_time.millis,
			stop_time.seconds, stop_time.millis,
			get_delta_time_in_ms(start_time, stop_time));
    // printf("InVoked: %lf seconds\n", sec(clock()-time));

    calc_network_cost_TA();
}
#define cluster_size (4)
#define OBJ_ID_1 "MASKMASK"
// extern char OBJ_ID_1[] = "MASKMASK";
#define OBJ_ID_SZ_1 8
#define OBJ_ID_2 "MASKMASKM"
// extern char OBJ_ID_2[] = "MASKMASKM";
#define OBJ_ID_SZ_2 9
// #define TEST_OBJECT_SIZE	(16 * 16 * 64)
#define TEST_OBJECT_SIZE	(1)
// #define 
extern int mask_chosen = 0;



TEE_Result create_raw_object(char* obj_id, size_t obj_id_sz, float *load_data, int size)
{
	TEE_ObjectHandle object;
	TEE_Result res;
	uint32_t obj_data_flag;
	/*
	 * Create object in secure storage and fill with data
	 */
	obj_data_flag = TEE_DATA_FLAG_ACCESS_READ |		/* we can later read the oject */
			TEE_DATA_FLAG_ACCESS_WRITE |		/* we can later write into the object */
			TEE_DATA_FLAG_ACCESS_WRITE_META |	/* we can later destroy or rename the object */
			TEE_DATA_FLAG_OVERWRITE;		/* destroy existing object of same ID */


    printf("tee_create_new_object\n");
	res = TEE_CreatePersistentObject(TEE_STORAGE_PRIVATE,
					obj_id, obj_id_sz,
					obj_data_flag,
					TEE_HANDLE_NULL,
					NULL, 0,		/* we may not fill it right now */
					&object);
    printf("tee_create_new_object1\n");
	if (res != TEE_SUCCESS) {
		EMSG("TEE_CreatePersistentObject failed 0x%08x", res);
		// TEE_Free(obj_id);
		return res;
	}
    printf("tee_create_new_object2\n");

	res = TEE_WriteObjectData(object, load_data, sizeof(float) * size);
	if (res != TEE_SUCCESS) {
		EMSG("TEE_WriteObjectData failed 0x%08x", res);
		TEE_CloseAndDeletePersistentObject1(object);
	} else {
		TEE_CloseObject(object);
	}
    printf("tee_create_new_object3\n");

	// TEE_Free(obj_id);
	return res;
}
TEE_Result read_raw_object(float *load_data, int size){
    TEE_ObjectHandle object;
	TEE_ObjectInfo object_info;
	TEE_Result res;
	uint32_t read_bytes;
    if (mask_chosen == 0){
        res = TEE_OpenPersistentObject(TEE_STORAGE_PRIVATE,
                    OBJ_ID_1, OBJ_ID_SZ_1,
                    TEE_DATA_FLAG_ACCESS_READ |
                    TEE_DATA_FLAG_SHARE_READ,
                    &object);        
    }
    else{
        res = TEE_OpenPersistentObject(TEE_STORAGE_PRIVATE,
            OBJ_ID_2, OBJ_ID_SZ_2,
            TEE_DATA_FLAG_ACCESS_READ |
            TEE_DATA_FLAG_SHARE_READ,
            &object);    
    }
    if (res != TEE_SUCCESS) {
		EMSG("Failed to open persistent object, res=0x%08x", res);
		return res;
	}
    res = TEE_GetObjectInfo1(object, &object_info);
	if (res != TEE_SUCCESS) {
		EMSG("Failed to create persistent object, res=0x%08x", res);
		return res;
	}
    res = TEE_ReadObjectData(object, load_data, sizeof(float) * size,
                &read_bytes);

	if (res != TEE_SUCCESS) {
		EMSG("TEE_ReadObjectData failed 0x%08x, read %" PRIu32 " over %u",
				res, read_bytes, object_info.dataSize);
        return res;
    }

}

void write_mask_test(float *obj1_data, int size){
    printf("[init_write_test]0\n");
    int x = 1;
    printf("[init_write_test]00\n");
    // float *obj1_data = malloc(sizeof(float) * size);
    printf("[init_write_test]1\n");
    for(int i = 0; i < TEST_OBJECT_SIZE; ++i) {
        obj1_data[i] = ta_rand();
    }
    printf("[init_write_test]2\n");

    create_raw_object(OBJ_ID_1, OBJ_ID_SZ_1, obj1_data, TEST_OBJECT_SIZE);
    printf("[init_write_test]3\n");
    for(int i = 0; i < TEST_OBJECT_SIZE; ++i) {
        obj1_data[i] = ta_rand();
    }
    printf("[init_write_test]4\n");
    create_raw_object(OBJ_ID_2, OBJ_ID_SZ_2, obj1_data, TEST_OBJECT_SIZE);
    printf("CREATE MASK FINISH\n");
    // free(obj1_data);
}

void tt_forward_relu(float *output, int size, int activation, int out_channels, int size_y, int use_mask){
    // TEE_Time start_time_mask = { };
	// TEE_Time stop_time_mask = { };
    printf("[tee]\n");
    int max_size_output_ = size;
	// TEE_GetSystemTime(&start_time_mask);
    float *temp_output_per_channel_y = malloc(sizeof(float) * size_y * size_y);
    // simulation load inv
    float *random_coeff_list_inv = malloc(sizeof(float) * out_channels * cluster_size);
    // for (int i = 0; i < out_channels * cluster_size; ++i){
    //     random_coeff_list_inv[i] = 0.01 * i;
    // }
    int * shuffle_idx = malloc(sizeof(int) * out_channels);
    read_raw_object(shuffle_idx, out_channels);
    read_raw_object(random_coeff_list_inv, out_channels * cluster_size);
	// TEE_GetSystemTime(&stop_time_mask);
    // printf("[RELU random fusion read]: %d(ms)\n",
	// 		get_delta_time_in_ms(start_time_mask, stop_time_mask));
    //    // simulation inv mastrix
	// TEE_GetSystemTime(&start_time_mask);

    for (int i = 0; i < out_channels / cluster_size; ++i){
        for(int j = 0; j < cluster_size; ++j) {
            for(int z = 0; z < size_y * size_y; ++z) {
                temp_output_per_channel_y[z] = 0;
                for (int k = 0; k  < cluster_size; ++k) {
                    temp_output_per_channel_y[z] += random_coeff_list_inv[i * cluster_size * cluster_size + j * cluster_size + k] * output[(i * cluster_size + k) * size_y * size_y + z];
                }
            }
            for (int z = 0; z < size_y * size_y; ++z) {
                output[(i * cluster_size + j) * size_y * size_y + z] = temp_output_per_channel_y[z];
            }
        }
    }
    // simulation deshuffle


    float *swap_temp_channel = malloc(sizeof(float) * max_size_output_);
    int k = 1;
    for (int i = 0; i < out_channels; ++i) {
        for (int j = 0; j < size_y * size_y; ++j){
            swap_temp_channel[((k + i) % out_channels) * size_y * size_y + j] = output[i * size_y * size_y + j];
        }
    }
    for (int i = 0; i < out_channels; ++i) {
        for (int j = 0; j < size_y * size_y; ++j){
            output[i * size_y * size_y + j] = swap_temp_channel[i * size_y * size_y + j];
        }
    }
	// TEE_GetSystemTime(&stop_time_mask);
    // printf("[RELU random fusion do]: %d(ms)\n",
	// 		get_delta_time_in_ms(start_time_mask, stop_time_mask));
    // // demask
	// TEE_GetSystemTime(&start_time_mask);
    float *load_data = malloc(sizeof(float) * max_size_output_);
    read_raw_object(load_data, max_size_output_);
	// TEE_GetSystemTime(&stop_time_mask);
    // printf("[RELU mask load]: %d(ms)\n",
	// 		get_delta_time_in_ms(start_time_mask, stop_time_mask));
	// TEE_GetSystemTime(&start_time_mask);
    if (use_mask > 0){
        for (int i = 0; i < max_size_output_; ++i) {
            output[i] -= load_data[i];
        }
    }
	// TEE_GetSystemTime(&stop_time_mask);
    // printf("[RELU mask del]: %d(ms)\n",
	// 		get_delta_time_in_ms(start_time_mask, stop_time_mask));
    // // relu
	// TEE_GetSystemTime(&start_time_mask);
    activate_array_TA(output, size, (ACTIVATION_TA) activation);
	// TEE_GetSystemTime(&stop_time_mask);
    // printf("[RELU mask relu]: %d(ms)\n",
	// 		get_delta_time_in_ms(start_time_mask, stop_time_mask));
    // add mask
	// TEE_GetSystemTime(&start_time_mask);
    mask_chosen = rand_int();
    read_raw_object(load_data, max_size_output_);
    for (int i = 0; i < max_size_output_; ++i) {
        output[i] += load_data[i];
    }
    // TEE_GetSystemTime(&stop_time_mask);
    // printf("[RELU mask add]:  %d(ms)\n",
	// 		get_delta_time_in_ms(start_time_mask, stop_time_mask));
    free(temp_output_per_channel_y);
    free(shuffle_idx);
    free(swap_temp_channel);
    free(load_data);
    // TEE_GetSystemTime(&stop_time);
    // printf("relu: start: %u.%u(s), stop: %u.%u(s), delta: %u(ms)",
	// 		start_time.seconds, start_time.millis,
	// 		stop_time.seconds, stop_time.millis,
	// 		get_delta_time_in_ms(start_time, stop_time));
}

void tt_forward_network_TA(int n_layer_th)
{
    TEE_Time start_time = { };
	TEE_Time stop_time = { };
	TEE_GetSystemTime(&start_time);

    netta.index = n_layer_th;
    layer_TA l = netta.layers[n_layer_th];

    if(l.delta){
        fill_cpu_TA(l.outputs * l.batch, 0, l.delta, 1);
    }

    l.forward_TA(l, netta);
    netta.input = l.output;
    // ta_net_input = malloc(sizeof(float)*l.outputs*l.batch);
    // for(int z=0; z<l.outputs*l.batch; z++){
    //     ta_net_input[z] = netta.input[z];
    // }
	TEE_GetSystemTime(&stop_time);

    printf("InVoked: start: %u.%u(s), stop: %u.%u(s), delta: %u(ms)",
			start_time.seconds, start_time.millis,
			stop_time.seconds, stop_time.millis,
			get_delta_time_in_ms(start_time, stop_time));
    
}

// 交换两行
void swapRows(float matrix[cluster_size][cluster_size * 2], int row1, int row2) {
    for (int i = 0; i < cluster_size * 2; i++) {
        float temp = matrix[row1][i];
        matrix[row1][i] = matrix[row2][i];
        matrix[row2][i] = temp;
    }
}

// 计算矩阵的逆
int invertMatrix(float *input, float inverse[cluster_size][cluster_size]) {
    float augmented[cluster_size][2 * cluster_size];
    
    // 初始化增广矩阵
    for (int i = 0; i < cluster_size; i++) {
        for (int j = 0; j < cluster_size; j++) {
            augmented[i][j] = *(input + i * cluster_size + j);
        }
        for (int j = cluster_size; j < 2 * cluster_size; j++) {
            augmented[i][j] = (i == (j - cluster_size)) ? 1.0 : 0.0;
        }
    }

    // 进行高斯-约旦消元法
    for (int i = 0; i < cluster_size; i++) {
        // 寻找主元
        if (augmented[i][i] == 0) {
            int swapRow = i;
            for (int j = i + 1; j < cluster_size; j++) {
                if (augmented[j][i] != 0) {
                    swapRow = j;
                    break;
                }
            }
            if (augmented[swapRow][i] == 0) {
                // 无法求逆
                return 0;
            }
            swapRows(augmented, i, swapRow);
        }

        // 归一化主元行
        float pivot = augmented[i][i];
        for (int j = 0; j < 2 * cluster_size; j++) {
            augmented[i][j] /= pivot;
        }

        // 消去其他行的元素
        for (int j = 0; j < cluster_size; j++) {
            if (j != i) {
                float factor = augmented[j][i];
                for (int k = 0; k < 2 * cluster_size; k++) {
                    augmented[j][k] -= factor * augmented[i][k];
                }
            }
        }
    }

    // 提取逆矩阵
    for (int i = 0; i < cluster_size; i++) {
        for (int j = 0; j < cluster_size; j++) {
            inverse[i][j] = augmented[i][j + cluster_size];
        }
    }

    return 1;
}

// obfuscation
void forward_network_TA_DEDefusion(int out_channels, int size_y)
{
    // float
    printf("out_channels: %d\n", out_channels);
    // out_channels = 16;
    if (out_channels % 4 > 0)
        return;
    float *random_coeff_list = malloc(sizeof(float) * out_channels * cluster_size);
    float *random_coeff_list_inv = malloc(sizeof(float) * out_channels * cluster_size);
    for(int i = 0; i < out_channels * cluster_size; ++i) {
        random_coeff_list[i] = (float)ta_rand_fusion();
        // printf("%lf ", random_coeff_list[i]);
    }
    float *pre_matrix;
    float inv_matrix[cluster_size][cluster_size];
    int ret = 0;
    
    // for (int i = 0; i < (int)(out_channels / cluster_size); ++i) {
    //     pre_matrix = &random_coeff_list[i * cluster_size * cluster_size];
    //     ret = invertMatrix(pre_matrix, inv_matrix);
    //     for(int j = 0; j < cluster_size * cluster_size; ++j){
    //         // printf("r: %d\n", j);
    //         random_coeff_list_inv[i * cluster_size * cluster_size + j] = *(&inv_matrix[0][0] + j);
    //     }
    // }

    float *output_y = malloc(sizeof(float) * out_channels * size_y * size_y);
    for (int i = 0; i < out_channels * size_y * size_y; ++i)
        output_y[i] = (float)ta_rand();
    printf("output_y[0]: %lf\n", output_y[0]);
    float *temp_output_per_channel_y = malloc(sizeof(float) * size_y * size_y);
    // defusion
    TEE_Time start_time_mask = { };
    TEE_Time middle_time_mask = { };
	TEE_Time stop_time_mask = { };
    float *output_y_mask = malloc(sizeof(float) * out_channels * size_y * size_y);
	TEE_GetSystemTime(&start_time_mask);
    // generate mask
    int MAX_TIMEs = 20;
    for (int j = 0; j < MAX_TIMEs; ++j){
        for (int i = 0; i < out_channels * size_y * size_y; ++i)
            output_y_mask[i] = (float)ta_rand();
    }
    // add mask 
	TEE_GetSystemTime(&stop_time_mask);
    
    printf("[MASK (RAND) COST]: start: %u.%u(s), stop: %u.%u(s), delta: %u(us)\n",
			start_time_mask.seconds, start_time_mask.millis,
			stop_time_mask.seconds, stop_time_mask.millis,
			get_delta_time_in_ms(start_time_mask, stop_time_mask) * 1000 / MAX_TIMEs);
	TEE_GetSystemTime(&start_time_mask);

    for (int j = 0; j < MAX_TIMEs; ++j) {
        for (int i = 0; i < out_channels * size_y * size_y; ++i)
            output_y[i] = output_y[i] + output_y_mask[i];
        // remove mask
        for (int i = 0; i < out_channels * size_y * size_y; ++i)
            output_y[i] = output_y[i] - output_y_mask[i];
    }
	TEE_GetSystemTime(&stop_time_mask);
    printf("[MASK (ADD + DEL) COST]: start: %u.%u(s), stop: %u.%u(s), delta: %u(us)\n",
			start_time_mask.seconds, start_time_mask.millis,
			stop_time_mask.seconds, stop_time_mask.millis,
			get_delta_time_in_ms(start_time_mask, stop_time_mask) * 1000 / MAX_TIMEs);
    TEE_Time start_time = { };
	TEE_Time stop_time = { };

    

    printf("start: out_channels: %d, cluster_size: %d, size_y: %d\n", out_channels, cluster_size, size_y);
	TEE_GetSystemTime(&start_time);
    // start_time = clock();
    MAX_TIMEs = 1000;
    for (int ii = 0; ii < MAX_TIMEs; ++ii) {
        for (int i = 0; i < out_channels / cluster_size; ++i){
            for(int j = 0; j < cluster_size; ++j) {
                for(int z = 0; z < size_y * size_y; ++z) {
                    temp_output_per_channel_y[z] = 0;
                    for (int k = 0; k  < cluster_size; ++k) {
                        // temp_output_per_channel_y[z] += random_coeff_list_inv[i * cluster_size * cluster_size + j * cluster_size + k] * output_y[(i * cluster_size + k) * size_y * size_y + z];
                        temp_output_per_channel_y[z] += random_coeff_list[i * cluster_size * cluster_size + j * cluster_size + k] * output_y[(i * cluster_size + k) * size_y * size_y + z];
                    }
                }
                for (int z = 0; z < size_y * size_y; ++z) {
                    output_y[(i * cluster_size + j) * size_y * size_y + z] = temp_output_per_channel_y[z];
                }
            }
        }
    }    
    // stop_time = clock();
	TEE_GetSystemTime(&stop_time);
    // printf("defusion: %lf(s)", sec(stop_time - start_time));
    printf("[Defusion]: start: %u.%u(s), stop: %u.%u(s), delta: %u(us)\n",
			start_time.seconds, start_time.millis,
			stop_time.seconds, stop_time.millis,
			get_delta_time_in_ms(start_time, stop_time) * 1000 / MAX_TIMEs);
    free(random_coeff_list);
    free(random_coeff_list_inv);
    free(output_y);
    free(temp_output_per_channel_y);
    free(output_y_mask);
    // free()
}


void forward_network_TA_TEST(int size)
{
    // if(roundnum == 0){
    //     // ta_net_input malloc so not destroy before addition backward
    //     ta_net_input = malloc(sizeof(float) * netta.layers[0].inputs * netta.layers[0].batch);
    //     ta_net_delta = malloc(sizeof(float) * netta.layers[0].inputs * netta.layers[0].batch);

    //     if(netta.workspace_size){
    //         printf("workspace_size=%ld\n", netta.workspace_size);
    //         netta.workspace = calloc(1, netta.workspace_size);
    //     }
    // }

    // roundnum++;
    // int i;
    // TEE_Time start_time = { };
	// TEE_Time stop_time = { };
	// TEE_GetSystemTime(&start_time);
    ta_net_input = malloc(sizeof(float) * size);
    // free(ta_net_input);

    return;
    // for(i = 0; i < netta.n; ++i){
    //     netta.index = i;
    //     layer_TA l = netta.layers[i];

    //     if(l.delta){
    //         fill_cpu_TA(l.outputs * l.batch, 0, l.delta, 1);
    //     }

    //     l.forward_TA(l, netta);

    //     // if(debug_summary_pass == 1){
    //     //     summary_array("forward_network / l.output", l.output, l.outputs*netta.batch);
    //     // }

    //     netta.input = l.output;

    //     if(l.truth) {
    //         netta.truth = l.output;
    //     }
    //     //output of the network (for predict)
    //     // &&
    //     // if(!netta.train && l.type == SOFTMAX_TA){
    //     //     ta_net_output = malloc(sizeof(float)*l.outputs*1);
    //     //     for(int z=0; z<l.outputs*1; z++){
    //     //         ta_net_output[z] = l.output[z];
    //     //     }
    //     // }

    //     if(i == netta.n - 1)  // ready to back REE for the rest forward pass
    //     {
    //         ta_net_input = malloc(sizeof(float)*l.outputs*l.batch);
    //         for(int z=0; z<l.outputs*l.batch; z++){
    //             ta_net_input[z] = netta.input[z];
    //         }
    //     }
    // }
	// TEE_GetSystemTime(&stop_time);
    // printf("InVoked: start: %u.%u(s), stop: %u.%u(s), delta: %u(ms)",
	// 		start_time.seconds, start_time.millis,
	// 		stop_time.seconds, stop_time.millis,
	// 		get_delta_time_in_ms(start_time, stop_time));
    // // printf("InVoked: %lf seconds\n", sec(clock()-time));

    // calc_network_cost_TA();
}


void update_network_TA(update_args_TA a)
{
    int i;
    for(i = 0; i < netta.n; ++i){
        layer_TA l = netta.layers[i];
        if(l.update_TA){
            l.update_TA(l, a);
        }
    }
}


void calc_network_cost_TA()
{
    int i;
    float sum = 0;
    int count = 0;
    for(i = 0; i < netta.n; ++i){
        if(netta.layers[i].cost){
            sum += netta.layers[i].cost[0];
            ++count;
        }
    }
    *netta.cost = sum/count;
    err_sum += *netta.cost;
}


void calc_network_loss_TA(int n, int batch)
{
    float loss = (float)err_sum/(n*batch);

    if(avg_loss == -1) avg_loss = loss;
    avg_loss = avg_loss*.9 + loss*.1;

    char loss_char[20];
    char avg_loss_char[20];
    ftoa(loss, loss_char, 5);
    ftoa(avg_loss, avg_loss_char, 5);
    IMSG("loss = %s, avg loss = %s from the TA\n",loss_char, avg_loss_char);
    err_sum = 0;
}



//void backward_network_TA(float *ca_net_input, float *ca_net_delta)
void backward_network_TA(float *ca_net_input)
{
    int i;

    for(i = netta.n-1; i >= 0; --i){
        layer_TA l = netta.layers[i];

        if(l.stopbackward) break;
        if(i == 0){
            for(int z=0; z<l.inputs*l.batch; z++){
             // note: both ca_net_input and ca_net_delta are pointer
                ta_net_input[z] = ca_net_input[z];
                //ta_net_delta[z] = ca_net_delta[z]; zeros removing
                ta_net_delta[z] = 0.0f;
            }

            netta.input = ta_net_input;
            netta.delta = ta_net_delta;
        }else{
            layer_TA prev = netta.layers[i-1];
            netta.input = prev.output;
            netta.delta = prev.delta;
        }

        netta.index = i;
        l.backward_TA(l, netta);

        // when the first layer in TEE is a Dropout layer
        if((l.type == DROPOUT_TA) && (i == 0)){
            for(int z=0; z<l.inputs*l.batch; z++){
                ta_net_input[z] = l.output[z];
                ta_net_delta[z] = l.delta[z];
            }
            //netta.input = l.output;
            //netta.delta = l.delta;
        }
    }
}
